use futures_util::{SinkExt, StreamExt};
use mimalloc::MiMalloc;
use rustls::{ClientConfig, RootCertStore};
use sonic_rs::{pointer, JsonValueTrait, Value};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio_tungstenite::tungstenite::{Message, Utf8Bytes};
use tokio_tungstenite::{connect_async_tls_with_config, Connector};
use webpki_roots::TLS_SERVER_ROOTS;

#[global_allocator]
static GLOBAL: MiMalloc = MiMalloc;

#[derive(Default)]
struct Ticker {
    /// 主所(买1, 卖1)
    left_exchange: (f64, f64),
    /// 副所(买1, 卖1)
    right_exchange: (f64, f64),
}

type Tickers = Arc<Mutex<Ticker>>;

#[tokio::main]
async fn main() {
    // TLS实现使用rustls
    rustls::crypto::ring::default_provider()
        .install_default()
        .expect("Failed to install rustls crypto provider");

    let ticker = Arc::new(Tickers::default());
    let ticker_clone_1 = Arc::clone(&ticker);
    let ticker_clone_2 = Arc::clone(&ticker);

    // 币安采集线程
    tokio::spawn(async move {
        // 加载根证书
        let root_store = RootCertStore {
            roots: TLS_SERVER_ROOTS.into(),
        };

        // 配置TLS文件
        let config = Arc::new(
            ClientConfig::builder()
                .with_root_certificates(root_store)
                .with_no_client_auth(),
        );

        loop {
            let connector = Connector::Rustls(Arc::clone(&config));

            // 连接币安
            let (ws_stream, _) = connect_async_tls_with_config(
                "wss://fstream.binance.com/ws/btcusdt@bookTicker",
                None,
                false,
                Some(connector),
            )
            .await
            .unwrap();

            // 切分流
            let (mut write, mut read) = ws_stream.split();

            loop {
                // 处理币安传来的消息
                match read.next().await.unwrap().unwrap() {
                    Message::Text(data) => {
                        // 将数据转换成document对象
                        // 解析到sonic_rs::Value后，底层是一个 Key-Value pair 的数组，而不会建立 HashMap 或 BTreeMap, 因此没有建表开销。
                        let root: Value =
                            unsafe { sonic_rs::from_slice_unchecked(data.as_bytes()).unwrap() };

                        let mut temp_ticker = ticker.lock().unwrap();

                        // 保存主所的买一卖一
                        temp_ticker.left_exchange = (
                            root.get("b").as_str().unwrap().parse().unwrap(),
                            root.get("a").as_str().unwrap().parse().unwrap(),
                        );
                    }
                    Message::Ping(_) => {
                        // 收到ping, 回复pong
                        write.send(Message::Pong(Default::default())).await.unwrap();
                    }
                    _ => break,
                }
            }
        }
    });

    // Bitget采集线程
    tokio::spawn(async move {
        // 加载根证书
        let root_store = RootCertStore {
            roots: TLS_SERVER_ROOTS.into(),
        };

        // 配置TLS文件
        let config = Arc::new(
            ClientConfig::builder()
                .with_root_certificates(root_store)
                .with_no_client_auth(),
        );

        loop {
            let connector = Connector::Rustls(Arc::clone(&config));

            // 连接bitget
            let (ws_stream, _) = connect_async_tls_with_config(
                "wss://ws.bitget.com/v2/ws/public",
                None,
                false,
                Some(connector),
            )
            .await
            .unwrap();

            // 切分流
            let (mut write, mut read) = ws_stream.split();

            // 订阅btc频道
            let data = sonic_rs::json!({
                "op": "subscribe",
                "args": [
                    {
                        "instType": "USDT-FUTURES",
                        "channel": "books1",
                        "instId": "BTCUSDT"
                    }
                ]
            });
            write
                .send(Message::Text(Utf8Bytes::from(data.to_string())))
                .await
                .unwrap();

            // 定时ping
            tokio::spawn(async move {
                loop {
                    tokio::time::sleep(Duration::from_secs(30)).await;

                    write
                        .send(Message::Text(Utf8Bytes::from("ping")))
                        .await
                        .unwrap();
                }
            });

            // 跳过 订阅成功回调
            read.next().await.unwrap().unwrap();

            loop {
                // 处理币安传来的消息
                while let Ok(data) = read.next().await.unwrap() {
                    // 跳过pong心跳
                    if data.len() == 4 {
                        continue;
                    }

                    if let Message::Text(data) = data {
                        // 这里我展示了另外一种提取数据的手段
                        // 优点是零拷贝
                        let bid: f64 = unsafe {
                            sonic_rs::get_from_slice_unchecked(
                                data.as_bytes(),
                                pointer!["data", 0, "bids", 0, 0],
                            )
                            .unwrap()
                        }
                        .as_str()
                        .unwrap()
                        .parse()
                        .unwrap();

                        let ask: f64 = unsafe {
                            sonic_rs::get_from_slice_unchecked(
                                data.as_bytes(),
                                pointer!["data", 0, "asks", 0, 0],
                            )
                            .unwrap()
                        }
                        .as_str()
                        .unwrap()
                        .parse()
                        .unwrap();

                        let mut temp_ticker = ticker_clone_2.lock().unwrap();

                        // 保存主所的买一卖一
                        temp_ticker.right_exchange = (bid, ask);

                        continue;
                    }

                    // 占位符, 不会走到这里
                    todo!()
                }
            }
        }
    });

    // 执行策略的线程
    let task = tokio::task::spawn_blocking(move || {
        loop {
            // 获取最新行情
            // 通过Drop trait来达到快速释放锁
            let temp_ticker = {
                let temp_ticker = ticker_clone_1.lock().unwrap();
                // 因为f64是Rust的基本类型,所以实现了Copy trait, 这里会执行按位复制,不会涉及堆分配
                (temp_ticker.left_exchange, temp_ticker.right_exchange)
            };

            // .......这里模拟执行策略
            println!("{:?}", temp_ticker);

            // 大概每10ms轮训一次新行情
            std::thread::sleep(std::time::Duration::from_millis(10));
        }
    });

    task.await.unwrap();
}
